{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "88b88f54",
   "metadata": {},
   "source": [
    "# MCP Setup and Integration with NeMo Agent Toolkit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e7b7dec",
   "metadata": {},
   "source": [
    "In this notebook we will walk through client and server support for [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) in the NVIDIA NeMo Agent Toolkit (NAT).\n",
    "\n",
    "Detailed documentation can be found in the [MCP Docs](../../docs/source/workflows/mcp/index.md), which inspired the execution flow of this notebook.\n",
    "\n",
    "**Goal**\n",
    "\n",
    "By the end of this notebook, you will be able to:\n",
    "- Use NAT agents as an **MCP server** to publish tools using MCP to be used by any MCP client.\n",
    "- Use NAT agents as an **MCP client** to connect to and use tools served by remote MCP servers.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d931d41f",
   "metadata": {},
   "source": [
    "## Table of Contents\n",
    "\n",
    "- [0.0) Setup](#setup)\n",
    "  - [0.1) Prerequisites](#prereqs)\n",
    "  - [0.2) API keys](#api-keys)\n",
    "  - [0.3) Installing NeMo Agent Toolkit](#install-nat)\n",
    "  - [0.4) Additional dependencies](#deps)\n",
    "- [1.0) NAT as an MCP Server](#mcp-servers)\n",
    "  - [1.1) Starting the server with default settings](#start-mcp-servers)\n",
    "  - [1.2) Starting the Server with optional flags](#start-mcp-servers-flags)\n",
    "  - [1.3) Filtering tools available from the server](#start-mcp-server-filter-tools)\n",
    "  - [1.4) Options for listing available tools](#server-listing-tool-options)\n",
    "- [2.0) NAT as an MCP client](#start-mcp-client)\n",
    "  - [2.1) Configure and run MCP clients using function groups](#mcp-client-function-groups)\n",
    "  - [2.2) Listing tools from MCP servers](#mcp-client-tool-calls)\n",
    "  - [2.3) Inspecting MCP client integration in workflows](#mcp-client-tools)\n",
    "  - [2.4) Calling individual MCP tools](#client-tools)\n",
    "- [3.0) Next steps](#next-steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa12603d",
   "metadata": {},
   "source": [
    "<a id=\"setup\"></a>\n",
    "# 0.0) Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4379db44",
   "metadata": {},
   "source": [
    "<a id=\"prereqs\"></a>\n",
    "## 0.1) Prerequisites\n",
    "\n",
    "We strongly recommend that users begin this notebook with a working understanding of NAT workflows. Please refer to earlier iterations of this notebook series prior to beginning this notebook.\n",
    "\n",
    "- **Platform:** Linux, macOS, or Windows\n",
    "- **Python:** version 3.11, 3.12, or 3.13\n",
    "- **Python Packages:** `pip`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dacfca5",
   "metadata": {},
   "source": [
    "<a id=\"api-keys\"></a>\n",
    "## 0.2) API Keys"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18bdce3c",
   "metadata": {},
   "source": [
    "For this notebook, you will need the following API keys to run all examples end-to-end:\n",
    "\n",
    "- **NVIDIA Build:** You can obtain an NVIDIA Build API Key by creating an [NVIDIA Build](https://build.nvidia.com) account and generating a key at https://build.nvidia.com/settings/api-keys\n",
    "\n",
    "Then you can run the cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a5ec8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "if \"NVIDIA_API_KEY\" not in os.environ:\n",
    "    nvidia_api_key = getpass.getpass(\"Enter your NVIDIA API key: \")\n",
    "    os.environ[\"NVIDIA_API_KEY\"] = nvidia_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a01741a7",
   "metadata": {},
   "source": [
    "<a id=\"install-nat\"></a>\n",
    "## 0.3) Installing NeMo Agent Toolkit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a4e0499",
   "metadata": {},
   "source": [
    "The recommended way to install NAT is through `pip` or `uv pip`.\n",
    "\n",
    "First, we will install `uv` which offers parallel downloads and faster dependency resolution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bd054b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "pip install uv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27aa8cde",
   "metadata": {},
   "source": [
    "NeMo Agent Toolkit can be installed through the PyPI `nvidia-nat` package.\n",
    "\n",
    "There are several optional subpackages available for NAT. For this example, we will rely on two subpackages:\n",
    "* The `nvidia-nat[langchain]` subpackage contains components for integrating with [LangChain](https://python.langchain.com/docs/introduction/).\n",
    "* The `nvidia-nat[mcp]` subpackage contains components for MCP clients and servers with NAT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc74f73b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "uv pip show -q \"nvidia-nat-langchain\"\n",
    "nat_langchain_installed=$?\n",
    "uv pip show -q \"nvidia-nat-mcp\"\n",
    "nat_mcp_installed=$?\n",
    "if [[ ${nat_langchain_installed} -ne 0 || ${nat_mcp_installed} -ne 0 ]]; then\n",
    "    uv pip install \"nvidia-nat[langchain,mcp]\"\n",
    "else\n",
    "    echo \"nvidia-nat[langchain,mcp] is already installed\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fd25bd7",
   "metadata": {},
   "source": [
    "<a id=\"deps\"></a>\n",
    "## 0.4) Additional dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d13dd4b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "uv pip show -q \"mcp-server-time\"\n",
    "mcp_time_installed=$?\n",
    "if [[ ${mcp_time_installed} -ne 0 ]]; then\n",
    "    uv pip install \"mcp-server-time\"\n",
    "else\n",
    "    echo \"mcp-server-time is already installed\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4de74aa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# install tomli-w to edit pyproject.toml and rename the package\n",
    "!uv pip install tomli tomli-w"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00419290",
   "metadata": {},
   "source": [
    "Next we will define a utility that will help us source the correct example code depending on your runtime environment and branch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c27df251",
   "metadata": {},
   "outputs": [],
   "source": [
    "def install_example_code(example_code_path: str):\n",
    "    \"\"\"\n",
    "    This function prompts the user to select an installation option for the example code.\n",
    "\n",
    "    Args:\n",
    "        example_code_path (str): The path to the example code to install.\n",
    "    \"\"\"\n",
    "    import shutil\n",
    "    import subprocess\n",
    "    from pathlib import Path\n",
    "\n",
    "    def rename_package_in_pyproject(pyproject_path: Path, suffix: str = \"_notebook\") -> str:\n",
    "        \"\"\"\n",
    "        Rename the package in pyproject.toml to avoid conflicts with existing installations.\n",
    "\n",
    "        Args:\n",
    "            pyproject_path: Path to the pyproject.toml file\n",
    "            suffix: Suffix to append to the package name\n",
    "\n",
    "        Returns:\n",
    "            The new package name, or None if renaming failed\n",
    "        \"\"\"\n",
    "        try:\n",
    "            # Try Python 3.11+ built-in tomllib first\n",
    "            try:\n",
    "                import tomllib as tomli\n",
    "            except ImportError:\n",
    "                import tomli\n",
    "            import tomli_w\n",
    "        except ImportError:\n",
    "            print(\"Warning: Could not import toml library, package name will not be modified\")\n",
    "            return None\n",
    "\n",
    "        try:\n",
    "            # Read the current pyproject.toml\n",
    "            with open(pyproject_path, 'rb') as f:\n",
    "                config = tomli.load(f)\n",
    "\n",
    "            # Get the original package name\n",
    "            original_name = config.get('project', {}).get('name')\n",
    "            if not original_name:\n",
    "                print(\"Warning: No package name found in pyproject.toml\")\n",
    "                return None\n",
    "\n",
    "            # Create new unique name\n",
    "            new_name = f\"{original_name}{suffix}\"\n",
    "\n",
    "            # Update the package name\n",
    "            config['project']['name'] = new_name\n",
    "\n",
    "            # Write back to pyproject.toml\n",
    "            with open(pyproject_path, 'wb') as f:\n",
    "                tomli_w.dump(config, f)\n",
    "\n",
    "            print(f\"✓ Renamed package from '{original_name}' to '{new_name}' to avoid conflicts\")\n",
    "            return new_name\n",
    "\n",
    "        except Exception as e:\n",
    "            print(f\"Warning: Could not modify package name: {e}\")\n",
    "            return None\n",
    "\n",
    "    # Detect if running in Colab\n",
    "    try:\n",
    "        import google.colab  # type: ignore  # noqa: F401\n",
    "    except ImportError:\n",
    "        in_colab = False\n",
    "        base_path = Path(os.getcwd())\n",
    "    else:\n",
    "        in_colab = True\n",
    "        base_path = Path(\"/content\")\n",
    "\n",
    "    # Simple input prompt for branch selection\n",
    "    print(\"=\" * 60)\n",
    "    print(\"Example code installation for: \", example_code_path)\n",
    "    print(\"=\" * 60)\n",
    "    print(\"\\nOptions:\")\n",
    "    if not in_colab:\n",
    "        print(\"  - Enter 'local' to copy example code to current directory\")\n",
    "    print(\"  - Enter a branch name (e.g., 'develop', 'main') to clone and install\")\n",
    "    print(\"=\" * 60)\n",
    "\n",
    "    from IPython.core.error import StdinNotImplementedError\n",
    "    try:\n",
    "        branch_name = input(\"\\nEnter your choice: \").strip()\n",
    "    except StdinNotImplementedError:\n",
    "        branch_name = 'local'\n",
    "        print(f\"\\nNo input available. Defaulting to branch: {branch_name}\")\n",
    "\n",
    "    # Default to 'main' if empty\n",
    "    if not branch_name:\n",
    "        branch_name = 'main'\n",
    "\n",
    "    if branch_name.lower() == 'local' and not in_colab:\n",
    "        # Copy files from local repository to current directory\n",
    "        print(\"\\nCopying example code from local repository...\")\n",
    "\n",
    "        # Try to find the repository root\n",
    "        current = Path.cwd()\n",
    "        repo_root = None\n",
    "\n",
    "        # Look for repository root by searching for pyproject.toml\n",
    "        for parent in [current] + list(current.parents):\n",
    "            if (parent / \"pyproject.toml\").exists() and (parent / \"src\" / \"nat\").exists():\n",
    "                repo_root = parent\n",
    "                break\n",
    "\n",
    "        if repo_root is None:\n",
    "            print(\"✗ Error: Could not find repository root.\")\n",
    "            print(\"Please use git install instead.\")\n",
    "            return\n",
    "\n",
    "        # path-check-skip-next-line\n",
    "        source_path = repo_root / example_code_path\n",
    "\n",
    "        # Get the example folder name (e.g., \"simple_calculator\")\n",
    "        example_name = Path(example_code_path).name\n",
    "        dest_path = base_path / f\"{example_name}_notebook\"\n",
    "\n",
    "        if source_path.exists():\n",
    "            # Remove destination if it already exists\n",
    "            if dest_path.exists():\n",
    "                print(f\"Removing existing directory: {dest_path}\")\n",
    "                shutil.rmtree(dest_path)\n",
    "\n",
    "            # Copy the entire example directory\n",
    "            shutil.copytree(source_path, dest_path)\n",
    "            print(f\"✓ Copied example code to: {dest_path.absolute()}\")\n",
    "\n",
    "            # Optional: Install dependencies if there's a pyproject.toml\n",
    "            pyproject_file = dest_path / \"pyproject.toml\"\n",
    "            if pyproject_file.exists():\n",
    "                # Rename package to avoid conflicts\n",
    "                new_name = rename_package_in_pyproject(pyproject_file)\n",
    "\n",
    "                print(\"\\nInstalling example package in editable mode...\")\n",
    "                get_ipython().system(f'pip install --no-deps -e {dest_path}')\n",
    "\n",
    "                if new_name:\n",
    "                    print(f\"✓ Installed package as '{new_name}' from: {dest_path}\")\n",
    "                else:\n",
    "                    print(f\"✓ Installed package from: {dest_path}\")\n",
    "\n",
    "            print(f\"\\nYou can now edit the files in: {dest_path}\")\n",
    "        else:\n",
    "            print(f\"✗ Error: Source path not found: {source_path.absolute()}\")\n",
    "            print(\"Make sure you're running this from the correct directory\")\n",
    "\n",
    "    elif branch_name.lower() == 'local' and in_colab:\n",
    "        print(\"\\n✗ Local install not available in Google Colab.\")\n",
    "        print(\"Please specify a branch name to clone from git.\")\n",
    "\n",
    "    else:\n",
    "        # Clone full repo and install from there\n",
    "        if in_colab:\n",
    "            print(f\"\\nCloning branch '{branch_name}' from GitHub...\")\n",
    "            repo_path = base_path / \"NeMo-Agent-Toolkit\"\n",
    "        else:\n",
    "            print(f\"\\nInstalling example code from branch: {branch_name}\")\n",
    "            repo_path = base_path / f\"NeMo-Agent-Toolkit-{branch_name}\"\n",
    "\n",
    "        repo_url = \"https://github.com/NVIDIA/NeMo-Agent-Toolkit.git\"\n",
    "\n",
    "        try:\n",
    "            # Remove existing repo if it exists\n",
    "            if repo_path.exists():\n",
    "                print(f\"Removing existing repository at: {repo_path}\")\n",
    "                shutil.rmtree(repo_path)\n",
    "\n",
    "            # Clone the full repository\n",
    "            print(f\"Cloning repository (branch: {branch_name})...\")\n",
    "            _ = subprocess.run(\n",
    "                [\"git\", \"clone\", \"--branch\", branch_name, repo_url, str(repo_path)],\n",
    "                capture_output=True,\n",
    "                text=True,\n",
    "                check=True\n",
    "            )\n",
    "            print(f\"✓ Cloned repository to: {repo_path}\")\n",
    "\n",
    "            # Verify the example path exists\n",
    "            example_path = repo_path / example_code_path\n",
    "            if not example_path.exists():\n",
    "                print(f\"✗ Error: Path '{example_code_path}' not found in branch '{branch_name}'\")\n",
    "                return\n",
    "\n",
    "            # Install the example package\n",
    "            pyproject_file = example_path / \"pyproject.toml\"\n",
    "            if pyproject_file.exists():\n",
    "                # Rename package to avoid conflicts\n",
    "                new_name = rename_package_in_pyproject(pyproject_file)\n",
    "\n",
    "                print(\"\\nInstalling example package in editable mode...\")\n",
    "                print(f\"Installing from: {example_path}\")\n",
    "                # get_ipython().system(f'pip install --no-deps -e {example_path}')\n",
    "                get_ipython().system(f'pip install --no-deps {example_path}')\n",
    "\n",
    "                if new_name:\n",
    "                    print(f\"✓ Installed package as '{new_name}'\")\n",
    "                else:\n",
    "                    print(\"✓ Installed package\")\n",
    "\n",
    "            # Get the example folder name\n",
    "            example_name = Path(example_code_path).name\n",
    "\n",
    "            print(\"\\n✓ Setup complete!\")\n",
    "            print(f\"Repository location: {repo_path}\")\n",
    "            print(f\"Example config: {example_path / 'configs'}\")\n",
    "            print(\"\\nYou can reference configs using:\")\n",
    "            print(f\"  {example_path / 'configs/config.yml'}\")\n",
    "\n",
    "            # Optionally copy configs to working directory for easy access\n",
    "            dest_path = base_path / f\"{example_name}_notebook\"\n",
    "            if dest_path.exists():\n",
    "                shutil.rmtree(dest_path)\n",
    "            shutil.copytree(example_path, dest_path)\n",
    "            print(f\"\\nAlso copied example to: {dest_path}\")\n",
    "            print(f\"Use configs from: {dest_path / 'configs/config.yml'}\")\n",
    "\n",
    "            # Remove existing repo if it exists\n",
    "            if repo_path.exists():\n",
    "                print(f\"Removing existing repository at: {repo_path}\")\n",
    "                shutil.rmtree(repo_path)\n",
    "\n",
    "\n",
    "        except subprocess.CalledProcessError as e:\n",
    "            print(f\"✗ Error cloning repository: {e.stderr}\")\n",
    "            print(f\"Make sure branch '{branch_name}' exists.\")\n",
    "        except Exception as e:\n",
    "            print(f\"✗ Error: {e}\")\n",
    "\n",
    "    print(\"\\n\" + \"=\" * 60)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aebcf9eb",
   "metadata": {},
   "source": [
    "<a id=\"mcp-servers\"></a>\n",
    "# 1.0) NAT as an MCP Server\n",
    "\n",
    "MCP uses a `client-server` architecture, where MCP clients connect to servers using one of a few transport types to discover and invoke tool calls from the distributed MCP servers. This core concept in agentic AI enables agents to determine which servers to invoke given the intent of the input prompt, and delegate tasks to their most appropriate service.\n",
    "\n",
    "Therefore, it's logical to start the tutorial of NAT support for MCP from the MCP server, where the new tools are hosted and called.\n",
    "\n",
    "The `nat mcp serve` command line utility can start an MCP server that publishes the functions from your workflow as MCP tools. In this section we are going to start up the `simple_calculator` server that we used to demonstrate the client-to-server connection in the first section, and dive deeper into hosting MCP servers with NAT."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01f44c88",
   "metadata": {},
   "source": [
    "<a id=\"start-mcp-servers\"></a>\n",
    "## 1.1) Starting the server with default settings\n",
    "\n",
    "First, let's install the `simple_calculator` example from NAT source as `simple_calculator_notebook`, which we can edit and run.\n",
    "\n",
    "We recommend using the 'develop' branch in Google Colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fe5946f",
   "metadata": {},
   "outputs": [],
   "source": [
    "install_example_code(\"examples/getting_started/simple_calculator\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a34822d",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash --bg\n",
    "# Then let's start up the simple calculator MCP server.\n",
    "nat mcp serve --config_file ./simple_calculator_notebook/configs/config.yml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8e1c22d",
   "metadata": {},
   "source": [
    "The `--bg` command will run your server in the background of the notebook, so you might not see confirmation that the server is running. To see the startup logs, you can remove this flag and run the cell again. However, subsequent steps require the server to be run as a background process.\n",
    "\n",
    "Next, we will list the set of available tools from `simple_calculator_notebook` MCP server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2d15d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait about 10 seconds before running the next cell...\n",
    "import time\n",
    "\n",
    "time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afdef7a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat mcp client tool list --url http://localhost:9901/mcp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91c13c22",
   "metadata": {},
   "source": [
    "The steps above have loaded the workflow configuration from the specified file, started an MCP server on the default host (localhost) and port (9901), and published all tools from the workflow as MCP tools. The MCP server is available at `http://localhost:9901/mcp` using streamable-http transport.\n",
    "\n",
    "You can also use the `sse` (Server-Sent Events) transport for backwards compatibility through the `--transport` flag, for example:\n",
    "```bash\n",
    "nat mcp serve --config_file ./simple_calculator_notebook/configs/config.yml --transport sse\n",
    "```\n",
    "With this configuration, the MCP server is available at `http://localhost:9901/sse` using SSE transport.\n",
    "\n",
    "<div style=\"color: red; font-style: italic;\">\n",
    "<strong>WARNING - SSE Transport Security Limitations</strong>: The SSE transport does not support authentication. For production deployments, use `streamable-http` transport with authentication configured. SSE should only be used for local development on localhost or behind an authenticating reverse proxy.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5429a4d",
   "metadata": {},
   "source": [
    "<a id=\"start-mcp-servers-flags\"></a>\n",
    "## 1.2) Starting the server with optional flags\n",
    "\n",
    "You can optionally specify the optional server runtime settings (host, port, process name) using the following flags, without directly editing the configuration file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1abe459e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's kill anything using port 9901 as it may be running from a previous section\n",
    "!lsof -ti:9901 | xargs kill -9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f6ef399",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash --bg\n",
    "nat mcp serve --config_file ./simple_calculator_notebook/configs/config.yml \\\n",
    "  --host 0.0.0.0 \\\n",
    "  --port 9901 \\\n",
    "  --name \"My MCP Server\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "111c81a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait about 10 seconds before running the next cell...\n",
    "import time\n",
    "\n",
    "time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ac1df26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For streamable-http transport (default)\n",
    "!nat mcp client tool list --url http://localhost:9901/mcp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17c2cc26",
   "metadata": {},
   "source": [
    "<a id=\"start-mcp-server-filter-tools\"></a>\n",
    "## 1.3) Filtering tools available from the server\n",
    "\n",
    "The `--tool_names` flag allows the user to specify a subset of tools that will be enabled on the server. In this example below, we will only enable `calculator.multiply` and `calculator.divide`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc8eb7d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's kill anything using port 9901 as it may be running from a previous section\n",
    "!lsof -ti:9901 | xargs kill -9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3e34b7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash --bg\n",
    "nat mcp serve --config_file ./simple_calculator_notebook/configs/config.yml \\\n",
    "  --tool_names calculator.multiply \\\n",
    "  --tool_names calculator.divide \\\n",
    "  --host 0.0.0.0 \\\n",
    "  --port 9901 \\\n",
    "  --name \"My MCP Server\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ce2eebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait about 10 seconds before running the next cell...\n",
    "import time\n",
    "\n",
    "time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38a9af39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For streamable-http transport (default)\n",
    "!nat mcp client tool list --url http://localhost:9901/mcp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21310a0f",
   "metadata": {},
   "source": [
    "Note in the above output that only `calculator.multiply` and `calculator.divide` are now available from the server.\n",
    "\n",
    "<a id=\"server-listing-tool-options\"></a>\n",
    "## 1.4) Options for listing available tools\n",
    "\n",
    "After running, you can verify the tool list available on that server using the NAT command line utilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "210fba97",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "nat mcp client tool list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07b514e7",
   "metadata": {},
   "source": [
    "Or by inspecting the exposed endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda8e515",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "curl -s http://localhost:9901/debug/tools/list | jq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f20d22b6",
   "metadata": {},
   "source": [
    "You can request one or more specific tools by name. The `name` parameter accepts repeated values or a comma‑separated list. When `name` is provided, detailed schemas are returned by default:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "721ee9a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Single tool (detailed by default)\n",
    "curl -s \"http://localhost:9901/debug/tools/list?name=calculator.multiply\" | jq\n",
    "\n",
    "# Multiple tools (detailed by default)\n",
    "curl -s \"http://localhost:9901/debug/tools/list?name=calculator.multiply&name=calculator.divide\" | jq\n",
    "\n",
    "# Comma-separated list (equivalent)\n",
    "curl -s \"http://localhost:9901/debug/tools/list?name=calculator.multiply,calculator.divide\" | jq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ac25dff",
   "metadata": {},
   "source": [
    "\n",
    "You can control the amount of detail using the `detail` query parameter:\n",
    "\n",
    "- When requesting specific tool(s) with `name`, detailed schema is returned by default. Pass `detail=false` to suppress schemas:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29fa8467",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "curl -s \"http://localhost:9901/debug/tools/list?name=calculator.multiply&detail=false\" | jq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3ca2c9f",
   "metadata": {},
   "source": [
    "- When listing all tools (without `name`), the default output is simplified. Pass `detail=true` to include schemas for each tool:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20b62f20",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "curl -s \"http://localhost:9901/debug/tools/list?detail=true\" | jq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba05bfad",
   "metadata": {},
   "source": [
    "<a id=\"start-mcp-client\"></a>\n",
    "# 2.0) NAT as an MCP client\n",
    "\n",
    "Using NAT to connect to MCP servers as a client is possible in two ways:\n",
    "1. `mcp_client`: A flexible configuration using function groups that allows you to connect to an MCP server, dynamically discover the tools it serves, and register them as NeMo Agent toolkit functions.\n",
    "2. `mcp_tool_wrapper`: A simple configuration that allows you to wrap a single MCP tool as a NeMo Agent toolkit function.\n",
    "\n",
    "In this section we are going to walk through use cases demonstrating both."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51602d83",
   "metadata": {},
   "source": [
    "<a id=\"mcp-client-function-groups\"></a>\n",
    "## 2.1) Configure and run MCP clients using function groups"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "412d2f2c",
   "metadata": {},
   "source": [
    "You can use the `mcp_client` function group to connect to an MCP server, dynamically discover the tools it serves, and register them as NeMo Agent toolkit functions. \n",
    "\n",
    "Key Features:\n",
    "- **Automatic tool discovery** from the MCP Server by the client.\n",
    "- **Easy YAML configurations**: The function group can be directly referenced in the workflow configuration. Multiple function groups can be used in the same workflow to access tools from multiple MCP servers. Refer to [Function Groups](../../docs/source/workflows/function-groups.md) for more information about function group capabilities.\n",
    "    - A tool within a function group can also be referenced by its name using the following syntax: `<function_group_name>.<tool_name>`.\n",
    "    - The function group supports filtering using the `include` and `exclude` parameters. You can also optionally override the tool name and description defined by the MCP server using the `tool_overrides` parameter. Note that `tool_overrides` require that the function is explicitly listed under the optional `include` list of the function group configuration. See [function group accessibility](../../docs/source/workflows/function-groups.md#understanding-function-accessibility) for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f41a42a",
   "metadata": {},
   "source": [
    "\n",
    "Here is a nice example of function group MCP client usage\n",
    "```\n",
    "function_groups:\n",
    "  mcp_tools:\n",
    "    _type: mcp_client\n",
    "    server:\n",
    "      transport: streamable-http\n",
    "      url: \"http://localhost:9901/mcp\"\n",
    "    include:\n",
    "      - tool_a\n",
    "      - tool_b\n",
    "    tool_overrides:\n",
    "      tool_a:\n",
    "        alias: \"tool_a_alias\"\n",
    "        description: \"Tool A description\"\n",
    "\n",
    "workflow:\n",
    "  _type: react_agent\n",
    "  tool_names:\n",
    "    - mcp_tools\n",
    "```\n",
    "\n",
    "We can also add individual tools using `mcp_tool_wrapper`\n",
    "```\n",
    "functions:\n",
    "  mcp_tool_a:\n",
    "    _type: mcp_tool_wrapper\n",
    "    url: \"http://localhost:9901/mcp\"\n",
    "    mcp_tool_name: tool_a\n",
    "  mcp_tool_b:\n",
    "    _type: mcp_tool_wrapper\n",
    "    url: \"http://localhost:9901/mcp\"\n",
    "    mcp_tool_name: tool_b\n",
    "\n",
    "workflows:\n",
    "  _type: react_agent\n",
    "  tool_names:\n",
    "    - mcp_tool_a\n",
    "    - mcp_tool_b\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a98570e7",
   "metadata": {},
   "source": [
    "If the server from Part 1 is still running, lets stop it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32bd151f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's kill anything using port 9901 before moving on to the next section for posterity\n",
    "!lsof -ti:9901 | xargs kill -9"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8624f234",
   "metadata": {},
   "source": [
    "Let's set up an MCP client using function groups.\n",
    "\n",
    "First we create a NAT workflow with the standard `nat workflow create` command line utility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52e4f754",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat workflow create mcp_dev_workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f158bf9f",
   "metadata": {},
   "source": [
    "Then let's add a new configuration for this MCP test. This configuration defines two function groups `mcp_time` and `mcp_math`.\n",
    "\n",
    "`mcp_time`: This `function_group` inherits from `_type: mcp_client`, which by default uses the streamable-http transport. However, we are overwriting this transport in lieu of the stdio transport, which requires us to define a startup command for the MCP server. This is necessary because the `mcp-server-time` package is designed to run as a local Python module rather than as a remote HTTP service. The stdio transport allows NAT to launch and communicate with the MCP server directly through standard input/output streams, making it ideal for locally-installed MCP servers. We specify the command `python -m mcp_server_time` to start the server process, along with any required arguments like the timezone configuration.\n",
    "\n",
    "`mcp_math`: This `function_group` also inherits from `_type: mcp_client`, but unlike `mcp_time`, it uses the default streamable-http transport. This transport type is designed for MCP servers that are already running as HTTP services. We configure it by specifying `url: \"http://localhost:9901/mcp\"`, which tells the NAT client where to connect to discover and invoke the available mathematical tools. The streamable-http transport is ideal for remote or containerized MCP servers that expose their functionality over HTTP endpoints, allowing for more flexible deployment scenarios compared to the stdio transport. However, this function group does require that the `mcp_math` server is up and running in advance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4544342",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./mcp_dev_workflow/configs/mcp_config.yml\n",
    "function_groups:\n",
    "  mcp_time:\n",
    "    _type: mcp_client\n",
    "    server:\n",
    "      transport: stdio\n",
    "      command: \"python\"\n",
    "      args: [\"-m\", \"mcp_server_time\", \"--local-timezone=America/Los_Angeles\"]\n",
    "  mcp_math:\n",
    "    _type: mcp_client\n",
    "    server:\n",
    "      transport: streamable-http\n",
    "      url: \"http://localhost:9901/mcp\"\n",
    "\n",
    "llms:\n",
    "  nim_llm:\n",
    "    _type: nim\n",
    "    model_name: meta/llama-3.1-70b-instruct\n",
    "    temperature: 0.0\n",
    "    max_tokens: 250\n",
    "\n",
    "workflow:\n",
    "  _type: react_agent\n",
    "  llm_name: nim_llm\n",
    "  tool_names:\n",
    "    - mcp_time\n",
    "    - mcp_math\n",
    "  verbose: true"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32afbcb0",
   "metadata": {},
   "source": [
    "Before running this workflow we need to start the `simple_calculator` server again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9af5a00",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash --bg\n",
    "nat mcp serve --config_file simple_calculator_notebook/configs/config.yml \\\n",
    "  --host 0.0.0.0 \\\n",
    "  --port 9901 \\\n",
    "  --name \"My MCP Server\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7efca58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait about 10 seconds before running the next cell...\n",
    "import time\n",
    "\n",
    "time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86d73e01",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nat run --config_file ./mcp_dev_workflow/configs/mcp_config.yml \\\n",
    "    --input \"Is the product of 2 * 4 greater than the current hour of the day in Los Angeles?\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c2a46cb",
   "metadata": {},
   "source": [
    "<a id=\"mcp-client-tool-calls\"></a>\n",
    "## 2.2) Listing tools from MCP servers\n",
    "\n",
    "Use the `nat mcp client` commands to inspect and call tools available from an MCP server before running your workflow. This is useful for discovering available tools from one or more servers and understanding their input schemas before integration with NAT workflows.\n",
    "\n",
    "First let's list the tools available from the `mcp_time` server. We should observe the following:\n",
    "\n",
    "```\n",
    "convert_time\n",
    "get_current_time\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fad15ee9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For stdio transport\n",
    "!nat mcp client tool list --transport stdio --command \"python\" --args \"-m mcp_server_time\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce50bd9d",
   "metadata": {},
   "source": [
    "Next let's list the tools available from the `mcp_math` server. We should observe the following:\n",
    "\n",
    "```\n",
    "calculator.subtract\n",
    "calculator.add\n",
    "current_datetime\n",
    "calculator.compare\n",
    "calculator.divide\n",
    "calculator.multiply\n",
    "react_agent\n",
    "````\n",
    "\n",
    "Note that the `react_agent` tool will be present in this list because it is a default chat capability that inherently gets deployed with the NAT server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac8978a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For streamable-http transport (default)\n",
    "!nat mcp client tool list --url http://localhost:9901/mcp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c9a191e",
   "metadata": {},
   "source": [
    "<a id=\"mcp-client-tools\"></a>\n",
    "## 2.3) Inspecting MCP client integration in workflows\n",
    "\n",
    "Once an MCP server has been integrated with a NAT workflow, it is often useful to inspect the tools integrated from the client side to confirm that tool discovery worked as expected.\n",
    "\n",
    "When you serve a workflow that includes an `mcp_client` function group, the NeMo Agent toolkit exposes an HTTP endpoint to inspect the tools configured on the client side and whether each tool is available on the connected server. We can inspect this workflow and as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3bc03b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's kill anything using port 9901 in case it's already running\n",
    "!lsof -ti:9901 | xargs kill -9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b00ef169",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash --bg\n",
    "nat mcp serve --config_file simple_calculator_notebook/configs/config.yml \\\n",
    "  --host 0.0.0.0 \\\n",
    "  --port 9901 \\\n",
    "  --name \"My MCP Server\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dc57c53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait about 10 seconds before running the next cell...\n",
    "import time\n",
    "\n",
    "time.sleep(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c34f7ed",
   "metadata": {},
   "source": [
    "<span style=\"color: red\"><i>note: If running this notebook in a cloud provider such as Google Colab, `dask` may be installed. If it is, you will first have to uninstall it via:</i></span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8df6918",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!uv pip uninstall dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22754f85",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's kill anything using port 8000 in case it's already running\n",
    "!lsof -ti:8000 | xargs kill -9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c8ffbd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash --bg\n",
    "nat serve --config_file ./mcp_dev_workflow/configs/mcp_config.yml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc1de7ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait about 10 seconds before running the next cell...\n",
    "import time\n",
    "\n",
    "time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29574ea8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Verify the workflow server is running\n",
    "!ps aux | grep \"nat serve\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bdc7303",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Call the endpoint and pretty-print the response\n",
    "!curl -s -v http://localhost:8000/mcp/client/tool/list | jq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ccffc2b",
   "metadata": {},
   "source": [
    "Example output from above:\n",
    "\n",
    "```\n",
    "{\n",
    "  \"mcp_clients\": [\n",
    "    {\n",
    "      \"function_group\": \"mcp_time\",\n",
    "      \"server\": \"stdio:python\",\n",
    "      \"transport\": \"stdio\",\n",
    "      \"session_healthy\": true,\n",
    "      \"protected\": false,\n",
    "      \"tools\": [\n",
    "        {\n",
    "          \"name\": \"convert_time\",\n",
    "          \"description\": \"Convert time between timezones\",\n",
    "          \"server\": \"stdio:python\",\n",
    "          \"available\": true\n",
    "        },\n",
    "        {\n",
    "          \"name\": \"get_current_time\",\n",
    "          \"description\": \"Get current time in a specific timezones\",\n",
    "          \"server\": \"stdio:python\",\n",
    "          \"available\": true\n",
    "        }\n",
    "      ],\n",
    "      \"total_tools\": 2,\n",
    "      \"available_tools\": 2\n",
    "    },\n",
    "    {\n",
    "      \"function_group\": \"mcp_math\",\n",
    "      \"server\": \"streamable-http:http://localhost:9901/mcp\",\n",
    "      \"transport\": \"streamable-http\",\n",
    "      \"session_healthy\": true,\n",
    "      \"protected\": false,\n",
    "      \"tools\": [\n",
    "        {\n",
    "          \"name\": \"calculator.add\",\n",
    "          \"description\": \"Add two or more numbers together.\",\n",
    "          \"server\": \"streamable-http:http://localhost:9901/mcp\",\n",
    "          \"available\": true\n",
    "        },\n",
    "        {\n",
    "          \"name\": \"calculator.compare\",\n",
    "          \"description\": \"Compare two numbers.\",\n",
    "          \"server\": \"streamable-http:http://localhost:9901/mcp\",\n",
    "          \"available\": true\n",
    "        },\n",
    "        {\n",
    "          \"name\": \"calculator.divide\",\n",
    "          \"description\": \"Divide one number by another.\",\n",
    "          \"server\": \"streamable-http:http://localhost:9901/mcp\",\n",
    "          \"available\": true\n",
    "        },\n",
    "        {\n",
    "          \"name\": \"calculator.multiply\",\n",
    "          \"description\": \"Multiply two or more numbers together.\",\n",
    "          \"server\": \"streamable-http:http://localhost:9901/mcp\",\n",
    "          \"available\": true\n",
    "        },\n",
    "        {\n",
    "          \"name\": \"calculator.subtract\",\n",
    "          \"description\": \"Subtract one number from another.\",\n",
    "          \"server\": \"streamable-http:http://localhost:9901/mcp\",\n",
    "          \"available\": true\n",
    "        },\n",
    "        {\n",
    "          \"name\": \"current_datetime\",\n",
    "          \"description\": \"Returns the current date and time in human readable format with timezone information.\",\n",
    "          \"server\": \"streamable-http:http://localhost:9901/mcp\",\n",
    "          \"available\": true\n",
    "        },\n",
    "        {\n",
    "          \"name\": \"react_agent\",\n",
    "          \"description\": \"ReAct Agent Workflow\",\n",
    "          \"server\": \"streamable-http:http://localhost:9901/mcp\",\n",
    "          \"available\": true\n",
    "        }\n",
    "      ],\n",
    "      \"total_tools\": 7,\n",
    "      \"available_tools\": 7\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e721f1f",
   "metadata": {},
   "source": [
    "<a id=\"client-tools\"></a>\n",
    "## 2.4) Calling individual MCP tools\n",
    "\n",
    "MCP tools can also be called individually using the `nat mcp client tool call` utility. This is very helpful for individual tool calls because it manages the MCP handshake and request simultaneously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f09d030b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Pass arguments as JSON\n",
    "nat mcp client tool call calculator.multiply \\\n",
    "  --url http://localhost:9901/mcp \\\n",
    "  --json-args '{\"numbers\": [2, 3]}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64b796a3",
   "metadata": {},
   "source": [
    "This section showed how to consume MCP servers by using NAT workflow command line utilities to deploy and manage MCP clients. With simple support for multiple transport types.\n",
    "\n",
    "NAT also supports [MCP Authentication](../../docs/source/workflows/mcp/mcp-auth.md).\n",
    "\n",
    "To use a protected MCP server, you need to provide the `--auth` flag:\n",
    "```bash\n",
    "nat mcp client tool list --url http://example.com/mcp --auth\n",
    "```\n",
    "This will use the `mcp_oauth2` authentication provider to authenticate the user. For more information, refer to [MCP Authentication](./mcp-auth.md).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28bd8876",
   "metadata": {},
   "source": [
    "<a id=\"next-steps\"></a>\n",
    "# 3.0 Next Steps\n",
    "\n",
    "MCP servers follow the MCP protocol, and therefore can be used by any MCP client, not only NAT MCP clients. As a follow-up exercise, try connecting an MCP server to your Agentic IDE like VS Code GitHub co-pilot, Cursor, Windsurf or Claude Code.\n",
    "\n",
    "Continue with the notebook series with `multi_agent_orchestration.ipynb`. Once agents are equipped with the tools they need to properly address user queries, we find a need to draw the line between one agent and another, and the concept of multi-agent systems and multi-agent orchestration takes center stage. In the next notebook, we will show how easily multi-agent orchestration can be coordinated using NAT and also show a quick demo of multi-modal asset handling with NAT agents."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "unew_312",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
